package controller

import (
	"chatgpt/ai"
	"chatgpt/ai/chatglm"
	openai2 "chatgpt/ai/openai"
	"chatgpt/helper"
	"chatgpt/helper/code"
	"chatgpt/services/completion"
	"github.com/gin-gonic/gin"
	"github.com/slclub/go-tips"
	"strings"
)

func QuestionHandle(ctx *gin.Context) {
	question := ctx.Query("q")
	mod := ctx.Query("model")
	//size := ctx.Query("size")
	//mod = "" // 屏蔽 图片生成
	//mcode, replay := openai2.RemoteAiGet(question, mod, size)
	//helper.Echo(ctx, int(mcode), replay)
	mcode, replay := ai.Completion(ctx, question, map[string]string{
		"model":  mod,
		"smodel": "chatglm_6b",
		"stops":  "", //completion.Mgr().StopSequence(uid),
		"unsafe": helper.UNSAFE_MOD,
		"stream": "",
	}, func() string {
		nquestion, _ := chatglm.Search.CustomQuestion(question)
		//if queue != nil {
		//	//nquestion = queue.ConvPrompt(question)
		//}
		return nquestion
	})
	helper.Echo(ctx, mcode, completion.ErrorCodeCompletion(tips.String(replay.Data), mcode), replay.Tokens)
}

func ImageGenerateHandle(ctx *gin.Context) {
	question := ctx.Query("q")
	size := ctx.Query("size")

	res := openai2.AppChatGPT.ConvImage(question, size)

	helper.Echo(ctx, int(res.MsgCode), res.Text)
}

func OriginHandle(ctx *gin.Context) {
	//question := ctx.Query("q")
	mod := ctx.Query("model")
	switch strings.Trim(mod, " ") {
	case "image-generate":
		ImageGenerateHandle(ctx)
		return
	case "chat":
		ChatHandle(ctx)
		return
	case "conversation":
		ConversationHandle(ctx)
		return
	case "completion", "":
		CompletionHandle(ctx)
		return
	}

	helper.Echo(ctx, int(100), "")
}

func CompletionHandle(ctx *gin.Context) {
	question := ctx.Query("q")
	mod := ctx.Query("model")
	//mod := ctx.Query("smodel")
	//unsafe := ctx.Query("unsafe")
	//stop := ctx.Query("stops")

	if !openai2.AppChatGPT.CheckQuestion(question) && !openai2.AppChatGPT.Unsafe(question, mod) {
		helper.Echo(ctx, int(code.SUCCESS), "抱歉没有找到与“"+question+"”相关的信息")
		return
	}

	//res := openai2.AppChatGPT.Conv(question, mod, stop)
	////if mcode == code.SUCCESS && !openai.AppChatGPT.Unsafe(question, mod) {
	////	replay = dirtywords.DirtyWords().Replace(replay)
	////}
	//res.Echo(ctx)
	mcode, replay := ai.Completion(ctx, question, map[string]string{
		"model":  mod,
		"smodel": "chatglm_6b",
		"stops":  "", //completion.Mgr().StopSequence(uid),
		"unsafe": helper.UNSAFE_MOD,
		"stream": "",
	}, func() string {
		nquestion, _ := chatglm.Search.CustomQuestion(question)
		//if queue != nil {
		//	//nquestion = queue.ConvPrompt(question)
		//}
		return nquestion
	})
	helper.Echo(ctx, mcode, completion.ErrorCodeCompletion(tips.String(replay.Data), mcode), replay.Tokens)
}

func ChatHandle(ctx *gin.Context) {
	question := ctx.Query("q")
	mod := ctx.Query("smodel")
	stop := ctx.Query("stops")
	uid := ctx.Query("uid")
	unsafe := ctx.Query("unsafe")

	if !openai2.AppChatGPT.CheckQuestion(question) && !openai2.AppChatGPT.Unsafe(question, unsafe) {
		helper.Echo(ctx, int(code.SUCCESS), "抱歉没有找到与“"+question+"”相关的信息")
		return
	}

	res := openai2.AppChatGPT.ConvChat(question, mod, &openai2.User{UID: uid}, stop)
	//if mcode == code.SUCCESS && !openai.AppChatGPT.Unsafe(question, mod) {
	//	replay = dirtywords.DirtyWords().Replace(replay)
	//}
	res.Echo(ctx)
}

func ConversationHandle(ctx *gin.Context) {
	question := ctx.Query("q")
	mod := ctx.Query("smodel")

	if !openai2.AppChatGPT.CheckQuestion(question) && !openai2.AppChatGPT.Unsafe(question, mod) {
		helper.Echo(ctx, int(code.SUCCESS), "抱歉没有找到与“"+question+"”相关的信息")
		return
	}

	res := openai2.AppChatGPT.ConvConversation(question, mod, nil)
	//if mcode == code.SUCCESS && !openai.AppChatGPT.Unsafe(question, mod) {
	//	replay = dirtywords.DirtyWords().Replace(replay)
	//}
	res.Echo(ctx)
}
